/*
 * Copyright 2025 The Authors
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *
 *@file    : imu_factor.h
 *@brief   : imu factor
 *@authors : zhanglei
 *@version : v1.0
 *@data    : 2025/9/14
 *
 */

#ifndef IMU_FACTOR_H
#define IMU_FACTOR_H

#include <ceres/ceres.h>

#include "preintegration/imu_preintegration.h"

#include "common/state.h"

namespace optimizer {

using namespace preintegration;

class ImuFactor : public ceres::SizedCostFunction<5, 1, 2, 2, 3, 1, 2, 2>{
public:  
   ImuFactor(std::shared_ptr<ImuPreIntegration> imu_preintegration);

   Eigen::Matrix<double, 5, 5> GetCov() {
      return imu_preintegration_->cov_;
   }
   // ceres-solver
   // parameters[0] : theta_i, vi, pi,  bg, ba;
   // parameters[1] : theta_j, vj, pj;
   virtual bool Evaluate(double const *const *parameters, double *residuals, double **jacobians) const {
      double Ri = parameters[0][0];
      Eigen::Vector2d vi(parameters[1][0], parameters[1][1]);

      Eigen::Vector2d pi(parameters[2][0], parameters[2][1]);

      double bg = parameters[3][0];
      Eigen::Vector2d ba(parameters[3][1], parameters[3][2]);
      
      double Rj = parameters[4][0]; 
      Eigen::Vector2d vj(parameters[5][0], parameters[5][1]);
  
      Eigen::Vector2d pj(parameters[6][0], parameters[6][1]);

      Eigen::Map<Eigen::Matrix<double, 5, 1>> residual(residuals);

      // LOG(INFO) << "vi " << vi.transpose() << std::endl;
      // LOG(INFO) << "pi " << pi.transpose() << std::endl;
      // LOG(INFO) << "ba " << ba.transpose() << std::endl;
      // LOG(INFO) << "vj " << vj.transpose() << std::endl;
      // LOG(INFO) << "pj " << pj.transpose() << std::endl;

      double dt = imu_preintegration_->dt_;

      SO2 dR = imu_preintegration_->GetDeltaRotation(bg);
      Eigen::Vector2d dv = imu_preintegration_->GetDeltaVelocity(bg, ba);
      Eigen::Vector2d dp = imu_preintegration_->GetDeltaPosition(bg, ba);

      Eigen::Matrix2d Rit = SO2(Ri).inverse(); 

      //（九）(30)
      residual(0,0) = SO2::exp(dR.inverse() * SO2(Ri).inverse() * SO2(Rj).matrix());

      //（九）(31)
      residual.block<2,1>(1,0) = Rit * (vj - vi) - dv;

      //（九）(32)
      residual.block<2,1>(3,0) = Rit * (pj - pi - vi * dt) - dp;
      
      // LOG(INFO) << "bg " << bg << ", ba : " << ba.transpose() << std::endl;
      // LOG(INFO) << "before residual " << Eigen::Matrix<double, 5, 1>(residual).transpose() << std::endl;
      // LOG(INFO) << "imu_preintegration_->cov_ \n " << imu_preintegration_->cov_ << std::endl;
      // LOG(INFO) << "imu_preintegration_->cov_.inverse() \n " << imu_preintegration_->cov_.inverse() << std::endl;

      Eigen::Matrix<double, 5, 5> sqrt_info = Eigen::LLT<Eigen::Matrix<double, 5, 5>>(imu_preintegration_->cov_.inverse()).matrixL().transpose();
      residual = sqrt_info * residual;

      // LOG(INFO) << "residual " << (residual).transpose() << std::endl;

      // LOG(INFO) << "imu sqrt_info \n " << sqrt_info << std::endl;

      //   col=0   col=1   col=2   col=3   col=4   col=5   col=6    col=7
      //   col=0   col=1,2 col=3,4 col=5   col=6,7 col=8   col=9,10 col=11,12
      // | dR1/dR1 dR1/dv1 dR1/dp1 dR1/dbg dR1/dba dR1/dR2 dR1/dv2 dR1/dp2| row=0, 0
      // | dv1/dR1 dv1/dv1 dv1/dp1 dv1/dbg dv1/dba dv1/dR2 dv1/dv2 dv1/dp2| row=1, 1~2
      // | dp1/dR1 dp1/dv1 dp1/dp1 dp1/dbg dp1/dba dp1/dR2 dp1/dv2 dp1/dp2| row=2, 3~4
      // | dR2/dR1 dR2/dv1 dR2/dp1 dR2/dbg dR2/dba dR2/dR2 dR2/dv2 dR2/dp2| row=3, 5
      // | dv2/dR1 dv2/dv1 dv2/dp1 dv2/dbg dv2/dba dv2/dR2 dv2/dv2 dv2/dp2| row=4, 6~7
      // | dp2/dR1 dp2/dv1 dp2/dp1 dp2/dbg dp2/dba dp2/dR2 dp2/dv2 dp2/dp2| row=5, 8~9
      if (jacobians) {
         if (jacobians[0]) { // dRi 
            Eigen::Map<Eigen::Matrix<double, 5, 1, Eigen::ColMajor>> jacobianRi(jacobians[0]);
            jacobianRi.setZero(); 
            // compute jacobian
            // dR1/dR1
            jacobianRi(0,0) = -1; // (33)
            // dv1/dbg
            jacobianRi.block<2,1>(1,0) = -SO2::so2(1) * Rit * (vj - vi); // (38)
            // dp1/dR1
            jacobianRi.block<2,1>(3,0) = -SO2::so2(1) * Rit * (pj - pi - vi * dt); // (42)
            jacobianRi = sqrt_info * jacobianRi;
         }
         if (jacobians[1]) { // dvi
            Eigen::Map<Eigen::Matrix<double, 5, 2, Eigen::RowMajor>> jacobianVi(jacobians[1]);
            jacobianVi.setZero(); 
            // dv1/dv1
            jacobianVi.block<2,2>(1,0) = -Rit; // (36)
            // dp1/dv1
            jacobianVi.block<2,2>(3,0) = -Rit * dt; // (41)
            jacobianVi = sqrt_info * jacobianVi;
         }
         if (jacobians[2]) { // dpi
            Eigen::Map<Eigen::Matrix<double, 5, 2, Eigen::RowMajor>> jacobianPi(jacobians[2]);
            jacobianPi.setZero();
            // dp1/dp1
            jacobianPi.block<2,2>(3,0) = -Rit; // (39)
            jacobianPi = sqrt_info * jacobianPi;
         }  
         if (jacobians[3]) { // dbg, dba
            Eigen::Map<Eigen::Matrix<double, 5, 3, Eigen::RowMajor>> jacobianBais(jacobians[3]);
            jacobianBais.setZero();
            // dR1/dbg1
            jacobianBais(0,0) = -dt; // (35)
            // dv1/dbg1
            jacobianBais.block<2,1>(1,0) = -imu_preintegration_->dV_dbg_; // -(27)
            // dv1/dba1
            jacobianBais.block<2,2>(1,1) = -imu_preintegration_->dV_dba_; // -(26)
            // dp/dbg1
            jacobianBais.block<2,1>(3,0) = -imu_preintegration_->dP_dbg_;
            // dp/dba1
            jacobianBais.block<2,2>(3,1) = -imu_preintegration_->dP_dba_;

            jacobianBais = sqrt_info * jacobianBais;
         } 

         if (jacobians[4]) { // dRj
            Eigen::Map<Eigen::Matrix<double, 5, 1, Eigen::ColMajor>> jacobianRj(jacobians[4]);
            jacobianRj.setZero();
            // dR/dR2
            jacobianRj(0,0) = 1; // (34)

            jacobianRj = sqrt_info * jacobianRj;
         }  

         if (jacobians[5]) { // dvj
            Eigen::Map<Eigen::Matrix<double, 5, 2, Eigen::RowMajor>> jacobianVj(jacobians[5]);
            jacobianVj.setZero();
            // dv/dv2
            jacobianVj.block<2,2>(1,0) = Rit; // (37)
            jacobianVj = sqrt_info * jacobianVj;
         }  

         if (jacobians[6]) {
            Eigen::Map<Eigen::Matrix<double, 5, 2, Eigen::RowMajor>> jacobianPj(jacobians[6]);
            jacobianPj.setZero();
            // dv/dv2
            jacobianPj.block<2,2>(3,0) = Rit; // (37)
            jacobianPj = sqrt_info * jacobianPj;
         }  
      }
   }

private:
   std::shared_ptr<ImuPreIntegration> imu_preintegration_;
};

}

#endif 
